package com.basker.pisces.springcloud.trace;

import java.lang.reflect.Method;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

import org.springframework.core.annotation.AnnotatedElementUtils;
import org.springframework.util.StringUtils;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.handler.HandlerInterceptorAdapter;

import com.basker.pisces.resource.StringResources;
import com.basker.pisces.springcloud.utils.RequestUtils;
import com.basker.pisces.trace.Span;
import com.basker.pisces.trace.Span.Tag;
import com.basker.pisces.trace.TraceContext;
import com.basker.pisces.trace.TraceContextHolder;
import com.basker.pisces.trace.annotation.EnableTrace;

import feign.Feign;

public class ServiceTraceContextInterceptor extends HandlerInterceptorAdapter {

    @Override
    public void afterCompletion(HttpServletRequest request, HttpServletResponse response, Object handler, Exception ex)
            throws Exception {
        TraceContext context = TraceContextHolder.getTraceContext();
        if (context == null) {
            return;
        }

        Span removed = context.removeCurrentSpan();
        if (removed != null) {
            context.publishSpanDestoryEvent(removed);
        }

        if (context.isEmpty()) {
            TraceContextHolder.remove();
        }
    }

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler)
            throws Exception {
        // 从消息头获取当前的Span
        Span spanFromRequest = RequestUtils.getSpanFromRequest();
        EnableTrace enableTrace = this.enableTrace(handler);

        if (spanFromRequest == null && enableTrace == null) {
            // 当前上下文没有span
            return true;
        }

        TraceContext context = new TraceContext();
        TraceContextHolder.setTraceContext(context, true);

        Span span = context.createRootSpan(spanFromRequest);

        if (enableTrace != null) {
            span.setTag(Tag.Server);
            span.setTraceEvents(enableTrace.events());
        } else {
            span.setTag(Tag.Noop);
        }

        span.setName(this.parseSpanName(handler, enableTrace));

        context.publishSpanCreateEvent(span);

        return true;
    }

    private EnableTrace enableTrace(Object handler) {
        if (!(handler instanceof HandlerMethod)) {
            return null;
        }

        HandlerMethod handlerMethod = (HandlerMethod) handler;
        Method method = handlerMethod.getMethod();
        return AnnotatedElementUtils.findMergedAnnotation(method, EnableTrace.class);
    }

    private String parseSpanName(Object handler, EnableTrace enableTrace) {
        HandlerMethod handlerMethod = (HandlerMethod) handler;
        Method method = handlerMethod.getMethod();

        if (StringUtils.isEmpty(enableTrace.action())) {
            return Feign.configKey(method.getDeclaringClass(), method);
        }

        return StringResources.resolve(enableTrace.action());
    }

}
